'''
Function:
    Define Depthwise Separable Convolution Module
Author:
    Zhenchao Jin
'''
# import torch
# import torch.nn as nn
import luojianet
import luojianet.nn as nn
import luojianet.ops as ops
from ..activation import BuildActivation
from ..normalization import BuildNormalization, constructnormcfg
# from activation import BuildActivation
# from normalization import BuildNormalization, constructnormcfg

'''DepthwiseSeparableConv2d'''
class DepthwiseSeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, bias=False, 
                 norm_cfg=None, act_cfg=None, dw_norm_cfg=None, dw_act_cfg=None, pw_norm_cfg=None, pw_act_cfg=None):
        super(DepthwiseSeparableConv2d, self).__init__()
        self.dw_norm_cfg = norm_cfg
        self.dw_act_cfg = act_cfg
        self.pw_norm_cfg = norm_cfg
        self.pw_act_cfg = act_cfg
        self.depthwise_conv = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=(padding[0], padding[0], padding[1], padding[1]) if isinstance(padding, tuple) else padding, dilation=dilation, group=in_channels, has_bias=bias, pad_mode='pad')
        if self.dw_norm_cfg is not None:
            self.depthwise_bn = BuildNormalization(constructnormcfg(placeholder=in_channels, norm_cfg=self.dw_norm_cfg))
        if self.dw_act_cfg is not None:
            self.depthwise_activate = BuildActivation(self.dw_act_cfg)
        self.pointwise_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, group=1, has_bias=bias, pad_mode='pad')
        if self.pw_norm_cfg is not None:
            self.pointwise_bn = BuildNormalization(constructnormcfg(placeholder=out_channels, norm_cfg=self.pw_norm_cfg))
        if self.pw_act_cfg is not None:
            self.pointwise_activate = BuildActivation(self.pw_act_cfg)
    '''forward'''
    def forward(self, x):
        x = self.depthwise_conv(x)
        if self.dw_norm_cfg is not None: x = self.depthwise_bn(x)
        if self.dw_act_cfg is not None: x = self.depthwise_activate(x)
        x = self.pointwise_conv(x)
        if self.pw_norm_cfg is not None: x = self.pointwise_bn(x)
        if self.pw_act_cfg is not None: x = self.pointwise_activate(x)
        return x